import tushare as ts
from datetime import datetime, timedelta
import pandas as pd
import functools
import time
import os

ts.set_token('2633a6ab6a5142792da59ea91f2204e3b8766bda519eeeae2df1c29c')  # 请替换为你的 token
pro = ts.pro_api()  # 初始化pro接口

# 保存下载后的数据
def save_dataframe_as_csv(df, filename_prefix, date_info = None, path="F:/Personal/data/"):
    if not date_info:
        # 获取当前日期并格式化为YYYYMMDD
        date_info = datetime.now().strftime('%Y%m%d')

    # 使用os.path.join构建完整的文件名（包括路径）
    filename = os.path.join(path, f"{filename_prefix}_{date_info}.csv")

    # 确保路径存在，不存在则创建
    os.makedirs(path, exist_ok=True)

    # 保存DataFrame为CSV文件
    df.to_csv(filename, index=False)  # index=False表示不保存DataFrame的索引

def get_ths_hot_data(start_date):
    # 将字符串转换为datetime对象

    df = pro.ths_hot(start_date=start_date)

    # 检查DataFrame是否为空（例如，没有行）
    if not df.empty:
        date_info = start_date

        file_prefix = "ths_hot"
        save_dataframe_as_csv(df, file_prefix, date_info)

        return df

    return None

# 计算并保存每天所有sec_code的
def get_combined_code_and_save(data_func, trade_date=None, filename_prefix='', path="F:/Personal/data/"):

    # 使用函数名作为文件名的一部分
    # Check if data_func is a functools.partial object
    if isinstance(data_func, functools.partial):
        func_name = data_func.args[0]
    else:
        func_name = data_func.__name__ if callable(data_func) else "no_func_name"

    # 初始化一个空的列表来存储结果
    data_list = []

    # 将字符串转换为datetime对象
    trade_dt = datetime.strptime(trade_date, '%Y%m%d')

    # 股票基础信息
    secu_df = pro.stock_basic(trade_date=trade_date)

    none_cnt = 0
    for index, row in secu_df.iterrows():
        ts_code = row['ts_code']
        # print(ts_code, trade_dt.strftime('%Y%m%d'))
        # 调用传入的数据源函数获取数据（日期范围）
        single_secu_data_df = data_func(ts_code=ts_code, trade_date=trade_dt.strftime('%Y%m%d'))
        if single_secu_data_df is None:
            none_cnt += 1
        if none_cnt >= 50:
            break
        time.sleep(0.2)

        data_list.append(single_secu_data_df)

    combined_df = pd.DataFrame()
    if data_list:

        # 将当天的所有证券信息何必为一个大DataFrame
        combined_df = pd.concat(data_list, ignore_index=True)

        # 设置文件名前缀
        filename_prefix = func_name

        # 构建日期信息
        date_info = trade_date

    # 检查DataFrame是否为空（例如，没有行）
    if combined_df.empty:
        print(f"{data_func} does not get data found on the {trade_date}.")
        return None

    # 保存DataFrame到CSV
    if not filename_prefix:
        filename_prefix = func_name

    save_dataframe_as_csv(combined_df, filename_prefix, date_info, path)

    return combined_df
def get_ths_member_data():

    # 同花顺概念板块成分
    df = pro.ths_member()

    file_prefix = "ths_member"
    save_dataframe_as_csv(df, file_prefix)

    return df

def get_hm_list_data():

    # 游资名录
    df = pro.hm_list()

    file_prefix = "hm_list"
    save_dataframe_as_csv(df, file_prefix)

    return df

def get_ths_index_data():

    # 同花顺概念和行业指数
    df = pro.ths_index()

    file_prefix = "ths_index"
    save_dataframe_as_csv(df, file_prefix)

    return df


def get_data_and_save(data_func, start_date=None, end_date=None, trade_date=None, month_date=None, filename_prefix='', path="F:/Personal/data/"):

    # 使用函数名作为文件名的一部分
    # Check if data_func is a functools.partial object
    if isinstance(data_func, functools.partial):
        func_name = data_func.args[0]
    else:
        func_name = data_func.__name__ if callable(data_func) else "no_func_name"

    if start_date and end_date:
        # 将字符串转换为datetime对象
        start_dt = datetime.strptime(start_date, '%Y%m%d')
        end_dt = datetime.strptime(end_date, '%Y%m%d')

        # 调用传入的数据源函数获取数据（日期范围）
        df = data_func(start_date=start_dt.strftime('%Y%m%d'), end_date=end_dt.strftime('%Y%m%d'))

        # 构建日期信息
        date_info = f"{start_date}_{end_date}"

    elif trade_date:
        # 将字符串转换为datetime对象
        trade_dt = datetime.strptime(trade_date, '%Y%m%d')

        # 调用传入的数据源函数获取数据（单日期）
        # 注意：这里假设数据源函数有一个参数可以接受单日期
        df = data_func(trade_date=trade_dt.strftime('%Y%m%d'))

        # 构建日期信息
        date_info = trade_date

    elif month_date:
        # 将字符串转换为datetime对象
        month_dt = datetime.strptime(month_date, '%Y%m%d')

        # 调用传入的数据源函数获取数据（单日期）
        # 注意：这里假设数据源函数有一个参数可以接受单日期
        df = data_func(month=month_dt.strftime('%Y%m%d')[:-2])

        # 构建日期信息
        date_info = month_date

    else:
        raise ValueError("Either start_date and end_date or trade_date must be provided.")

        # 检查DataFrame是否为空（例如，没有行）
    if df.empty and start_date and end_date:
        print(f"{data_func} does not get data found for the {start_date} and {end_date} range.")
        return None
    elif df.empty and trade_date:
        print(f"{data_func} does not get data found on the {trade_date}.")
        return None

    # 保存DataFrame到CSV
    if not filename_prefix:
        filename_prefix = func_name
    save_dataframe_as_csv(df, filename_prefix, date_info, path)

    return df

def get_daily_loop_data(start_date, end_date):

    # 将字符串转换为datetime对象
    start_dt = datetime.strptime(start_date, '%Y%m%d')
    end_dt = datetime.strptime(end_date, '%Y%m%d')

    df = pd.DataFrame()
    # 使用循环遍历日期
    current_dt = start_dt
    last_month_date = None
    while current_dt <= end_dt:
        is_weekend = current_dt.weekday() in [5, 6]

        if is_weekend:
            current_dt += timedelta(days=1)
            continue

        time.sleep(0.2)

        # 将datetime对象转换为'YYYYMMDD'格式的字符串
        trade_date = current_dt.strftime('%Y%m%d')
        #
        # 个股资金流向
        get_data_and_save(pro.moneyflow, trade_date=trade_date)
        #
        # # 龙虎榜每日明细
        # get_data_and_save(pro.top_list, trade_date=trade_date)
        #
        # # 龙虎榜机构明细
        # get_data_and_save(pro.top_inst, trade_date=trade_date)
        #
        # # # 卖方盈利预测数据
        # # get_data_and_save(pro.report_rc, trade_date = trade_date)
        #
        # # 每日筹码及胜率
        # if current_dt >= datetime.strptime('20180102', '%Y%m%d'):
        #     get_data_and_save(pro.cyq_perf, trade_date=trade_date)
        #
        # # 每日筹码分布
        # # start_time = time.time()
        # # get_combined_code_and_save(pro.cyq_chips, trade_date=trade_date)
        # # end_time = time.time()
        # # run_time = end_time - start_time
        # # print(f"每日筹码分布，单日运行时间{run_time}秒")
        #
        # # 股票技术因子（量化因子）
        # get_data_and_save(pro.stk_factor, trade_date=trade_date)
        #
        # if current_dt >= datetime.strptime('20170101', '%Y%m%d'):
        #     # 沪深港股通持股明细
        #     get_data_and_save(pro.hk_hold, trade_date=trade_date)
        #     # 沪深港通资金流向
        #     get_data_and_save(pro.moneyflow_hsgt, trade_date=trade_date)

        #
        # if current_dt >= datetime.strptime('20200101', '%Y%m%d'):
        #     # 涨跌停列表（新）
        #     get_data_and_save(pro.limit_list_d, trade_date=trade_date)

        # 券商每月荐股
        #
        # month_date = current_dt.strftime('%Y%m%d')
        # if current_dt >= datetime.strptime('20210101', '%Y%m%d'):
        #     if last_month_date != month_date:
        #         get_data_and_save(pro.broker_recommend, month_date=month_date)
        #         last_month_date = current_dt.strftime('%Y%m%d')
        #
        #
        #
        # # 大盘指数每日指标
        # get_data_and_save(pro.index_dailybasic, trade_date=trade_date)
        #
        # # 市场交易统计
        # get_data_and_save(pro.daily_info, trade_date=trade_date)
        #
        # # 深圳市场每日交易概况
        # get_data_and_save(pro.sz_daily_info, trade_date=trade_date)
        #
        # # 同花顺板块指数行情
        # get_data_and_save(pro.ths_daily.sql, trade_date=trade_date)
        #
        # # 游资每日明细
        # if current_dt >= datetime.strptime('20220801', '%Y%m%d'):
        #     get_data_and_save(pro.hm_detail, trade_date=trade_date)
        #
        # # 同花顺热榜
        # if current_dt >= datetime.strptime('20240101', '%Y%m%d'):
        #     get_data_and_save(pro.ths_hot, trade_date=trade_date)

        current_dt += timedelta(days=1)



# 调用函数并获取数据
get_daily_loop_data('20100101', '20240524')
# get_ths_member_data()
# get_hm_list_data()
# get_ths_index_data()
# get_ths_hot_data('20230101')